package com.liss.lissaiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class FileBasedChatMemory implements ChatMemory {

    private final String BASE_DIR;
    private static final Kryo kryo = new Kryo();

    static {
        // 关闭强制类注册，避免未注册类报错
        kryo.setRegistrationRequired(false);
        // 解决无参构造缺失的类实例化问题（适配 Spring AI 消息类）
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public FileBasedChatMemory(String dir) {
        this.BASE_DIR = dir;
        File baseDir = new File(dir);
        // 确保存储目录存在，不存在则创建
        if (!baseDir.exists()) {
            baseDir.mkdirs();
        }
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        // 获取已有对话，追加新消息后保存
        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        conversationMessages.addAll(messages);
        saveConversation(conversationId, conversationMessages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> allMessages = getOrCreateConversation(conversationId);
        // 截取最后 N 条消息，若消息总数不足 N 则返回全部
        int startIndex = Math.max(0, allMessages.size() - lastN);
        return allMessages.stream()
                .skip(startIndex)
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        // 删除对应对话文件，实现“清空记忆”
        File file = getConversationFile(conversationId);
        if (file.exists()) {
            file.delete();
        }
    }

    /**
     * 获取已有对话，若不存在则返回空列表
     */
    private List<Message> getOrCreateConversation(String conversationId) {
        File file = getConversationFile(conversationId);
        List<Message> messages = new ArrayList<>();

        // 若文件存在，反序列化读取对话
        if (file.exists()) {
            try (Input input = new Input(new FileInputStream(file))) {
                messages = kryo.readObject(input, ArrayList.class);
            } catch (IOException e) {
                System.err.println("读取对话文件失败（conversationId：" + conversationId + "）");
                e.printStackTrace();
            }
        }

        return messages;
    }

    /**
     * 将对话列表序列化保存到文件
     */
    private void saveConversation(String conversationId, List<Message> messages) {
        File file = getConversationFile(conversationId);
        // 用 try-with-resources 自动关闭流，避免资源泄漏
        try (Output output = new Output(new FileOutputStream(file))) {
            kryo.writeObject(output, messages);
        } catch (IOException e) {
            System.err.println("保存对话文件失败（conversationId：" + conversationId + "）");
            e.printStackTrace();
        }
    }

    /**
     * 生成对话对应的文件（每个对话一个独立 .kryo 文件）
     */
    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".kryo");
    }
}